import pandas as pd
import numpy as np
from scipy.stats import pointbiserialr
import matplotlib.pyplot as plt


def calculate_point_biserial_correlation(df, col1, col2):
    """
    计算 DataFrame 中两个指定列之间的点双列相关系数和 p 值。

    参数:
    df (DataFrame): 包含数据的 DataFrame。
    col1 (str): 第一个变量的列名，将被转换为二元变量。
    col2 (str): 第二个变量的列名（连续数值）。

    返回:
    point_biserial_corr (float): 点双列相关系数。
    p_value (float): p 值。
    """
    # 提取列
    col1_data = df[col1]
    col2_data = df[col2]

    # 将col1转换为二元变量
    col1_binary = np.where(col1_data > 0, 1, -1)

    # 计算点双列相关系数
    point_biserial_corr, p_value = pointbiserialr(col1_binary, col2_data)

    # 打印结果
    print(f"{col1} (binary) 和 {col2} 之间的点双列相关系数: {point_biserial_corr}")
    print(f"p值: {p_value}")

    # 可视化
    plt.figure(figsize=(10, 6))
    plt.scatter(col1_binary, col2_data, alpha=0.5)
    plt.xlabel(f'{col1} (binary)')
    plt.ylabel(col2)
    plt.title(f'{col1} (binary) vs {col2}')
    plt.show()

    return point_biserial_corr, p_value


# 示例数据
np.random.seed(42)  # 确保可重复性
size = 10  # 样本大小
data = {
    'col1': np.random.uniform(-0.1, 0.1, size),
    'col2': np.random.uniform(0, 100, size)
}
df = pd.DataFrame(data)

# 调用函数
calculate_point_biserial_correlation(df, 'col1', 'col2')
